"""
模型训练模块
负责模型选择、训练和调优
"""

import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV, cross_val_score
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import joblib
import os


def train_logistic_regression(X_train, y_train, random_state=42):
    """
    训练逻辑回归模型

    参数:
        X_train (pd.DataFrame): 训练集特征
        y_train (pd.Series): 训练集标签
        random_state (int): 随机种子

    返回:
        LogisticRegression: 训练好的逻辑回归模型
    """
    print("训练逻辑回归模型...")

    # 创建逻辑回归模型
    lr_model = LogisticRegression(random_state=random_state, max_iter=1000)

    # 训练模型
    lr_model.fit(X_train, y_train)

    print("逻辑回归模型训练完成")

    return lr_model


def train_random_forest(X_train, y_train, random_state=42):
    """
    训练随机森林模型

    参数:
        X_train (pd.DataFrame): 训练集特征
        y_train (pd.Series): 训练集标签
        random_state (int): 随机种子

    返回:
        RandomForestClassifier: 训练好的随机森林模型
    """
    print("训练随机森林模型...")

    # 创建随机森林模型
    rf_model = RandomForestClassifier(random_state=random_state)

    # 训练模型
    rf_model.fit(X_train, y_train)

    print("随机森林模型训练完成")

    return rf_model


def optimize_random_forest(X_train, y_train, random_state=42):
    """
    使用网格搜索优化随机森林模型

    参数:
        X_train (pd.DataFrame): 训练集特征
        y_train (pd.Series): 训练集标签
        random_state (int): 随机种子

    返回:
        tuple: (优化后的模型, 最佳参数)
    """
    print("优化随机森林模型...")

    # 定义参数网格
    param_grid = {
        'n_estimators': [50, 100, 200],
        'max_depth': [None, 10, 20, 30],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [1, 2, 4]
    }

    # 创建随机森林模型
    rf_model = RandomForestClassifier(random_state=random_state)

    # 创建网格搜索对象
    grid_search = GridSearchCV(
        estimator=rf_model,
        param_grid=param_grid,
        cv=5,
        scoring='accuracy',
        n_jobs=-1,
        verbose=1
    )

    # 执行网格搜索
    grid_search.fit(X_train, y_train)

    # 获取最佳模型和参数
    best_rf = grid_search.best_estimator_
    best_params = grid_search.best_params_

    print(f"随机森林模型优化完成")
    print(f"最佳参数: {best_params}")
    print(f"最佳交叉验证分数: {grid_search.best_score_:.4f}")

    return best_rf, best_params


def evaluate_model_cv(model, X, y, cv=5):
    """
    使用交叉验证评估模型

    参数:
        model: 要评估的模型
        X (pd.DataFrame): 特征数据
        y (pd.Series): 标签数据
        cv (int): 交叉验证折数

    返回:
        dict: 包含各项评估指标的字典
    """
    print(f"执行{cv}折交叉验证...")

    # 计算各项指标
    accuracy_scores = cross_val_score(model, X, y, cv=cv, scoring='accuracy')
    precision_scores = cross_val_score(model, X, y, cv=cv, scoring='precision')
    recall_scores = cross_val_score(model, X, y, cv=cv, scoring='recall')
    f1_scores = cross_val_score(model, X, y, cv=cv, scoring='f1')

    # 计算平均值和标准差
    metrics = {
        'accuracy': {
            'mean': accuracy_scores.mean(),
            'std': accuracy_scores.std()
        },
        'precision': {
            'mean': precision_scores.mean(),
            'std': precision_scores.std()
        },
        'recall': {
            'mean': recall_scores.mean(),
            'std': recall_scores.std()
        },
        'f1': {
            'mean': f1_scores.mean(),
            'std': f1_scores.std()
        }
    }

    print(f"交叉验证完成")
    print(f"准确率: {metrics['accuracy']['mean']:.4f} (±{metrics['accuracy']['std']:.4f})")
    print(f"精确率: {metrics['precision']['mean']:.4f} (±{metrics['precision']['std']:.4f})")
    print(f"召回率: {metrics['recall']['mean']:.4f} (±{metrics['recall']['std']:.4f})")
    print(f"F1分数: {metrics['f1']['mean']:.4f} (±{metrics['f1']['std']:.4f})")

    return metrics


def save_model(model, filepath):
    """
    保存训练好的模型

    参数:
        model: 要保存的模型
        filepath (str): 模型保存路径
    """
    # 确保目录存在
    os.makedirs(os.path.dirname(filepath), exist_ok=True)

    # 保存模型
    joblib.dump(model, filepath)

    print(f"模型已保存到: {filepath}")


def load_model(filepath):
    """
    加载保存的模型

    参数:
        filepath (str): 模型文件路径

    返回:
        加载的模型
    """
    model = joblib.load(filepath)
    print(f"模型已从 {filepath} 加载")
    return model


def train_and_select_best_model(X_train, y_train, X_val=None, y_val=None, optimize=True, random_state=42):
    """
    训练多个模型并选择最佳模型

    参数:
        X_train (pd.DataFrame): 训练集特征
        y_train (pd.Series): 训练集标签
        X_val (pd.DataFrame, optional): 验证集特征
        y_val (pd.Series, optional): 验证集标签
        optimize (bool): 是否优化随机森林模型
        random_state (int): 随机种子

    返回:
        tuple: (最佳模型, 模型评估结果)
    """
    models = {}
    evaluations = {}

    # 1. 训练逻辑回归模型
    lr_model = train_logistic_regression(X_train, y_train, random_state)
    models['Logistic Regression'] = lr_model

    # 评估逻辑回归模型
    lr_cv_metrics = evaluate_model_cv(lr_model, X_train, y_train)
    evaluations['Logistic Regression'] = lr_cv_metrics

    # 2. 训练随机森林模型
    if optimize:
        rf_model, best_params = optimize_random_forest(X_train, y_train, random_state)
    else:
        rf_model = train_random_forest(X_train, y_train, random_state)

    models['Random Forest'] = rf_model

    # 评估随机森林模型
    rf_cv_metrics = evaluate_model_cv(rf_model, X_train, y_train)
    evaluations['Random Forest'] = rf_cv_metrics

    # 如果有验证集，在验证集上评估
    if X_val is not None and y_val is not None:
        print("在验证集上评估模型...")

        for name, model in models.items():
            y_pred = model.predict(X_val)

            val_accuracy = accuracy_score(y_val, y_pred)
            val_precision = precision_score(y_val, y_pred)
            val_recall = recall_score(y_val, y_pred)
            val_f1 = f1_score(y_val, y_pred)

            evaluations[name]['validation'] = {
                'accuracy': val_accuracy,
                'precision': val_precision,
                'recall': val_recall,
                'f1': val_f1
            }

            print(f"{name} 验证集结果:")
            print(f"准确率: {val_accuracy:.4f}")
            print(f"精确率: {val_precision:.4f}")
            print(f"召回率: {val_recall:.4f}")
            print(f"F1分数: {val_f1:.4f}")

    # 选择最佳模型（基于交叉验证的F1分数）
    best_model_name = max(evaluations.keys(), key=lambda x: evaluations[x]['f1']['mean'])
    best_model = models[best_model_name]

    print(f"最佳模型: {best_model_name}")
    print(f"交叉验证F1分数: {evaluations[best_model_name]['f1']['mean']:.4f}")

    return best_model, evaluations
